'''
gpu269.mistgpu.com:10004
nohup python -u 04_Gru2Grutrain.py > ./output/04_Gru2Grutrain.txt 2>&1 &
nohup tensorboard --logdir=./tensorboard --port=10004 --host=0.0.0.0 > ./output/tensorboard.txt 2>&1 &
tail -f output/04_Gru2Grutrain.txt
'''
import os
import time
import tensorflow as tf
import numpy as np
from tensorflow.python.keras.callbacks import ModelCheckpoint

from ShipLoader import ShipLoader
# import Model
import os
from models import create_gru2gru
from utils.metrics import RMSE, root_mean_squared_error

file_path = "./DataSet/datasets/liner_npy"
# parameters for traning
learnig_rate = 0.001
epochs = 50
batch_size = 32

# parameters for seq2seq model
n_lstm = 128
encoder_length = 40
decoder_length = 20
now_time = time.strftime('%Y-%m-%d_%H-%M-%S',time.localtime(time.time()))
settings = '{}_{}_{}_{}'.format('gru2gru', epochs, batch_size, now_time)

# Load Ship data.
seq2seq_loader = ShipLoader(file_path, 'train')
seq2seq_loader.loadShipData()
train_x, train_y = seq2seq_loader.get_all_data(flag='train')

seq2seq_loader = ShipLoader(file_path, 'valid')
seq2seq_loader.loadShipData()
test_x, test_y = seq2seq_loader.get_all_data(flag='valid')
# 创建model
gru2gru_model = create_gru2gru(encoder_length, n_lstm, decoder_length)
# Choose Adam optimizer.
optimizer = tf.keras.optimizers.Adam(learnig_rate)
# seq2seq_model.compile(loss='rmse', metrics='mse', optimizer=optimizer)
# loss=tf.keras.metrics.mean_squared_error,
# metrics=[tf.keras.metrics.RootMeanSquaredError(name='rmse')]
gru2gru_model.compile(loss=root_mean_squared_error, metrics=['mse',tf.keras.metrics.RootMeanSquaredError(name='rmse'),'mae'], optimizer=optimizer)
gru2gru_model.summary()

callback_lists = []
callback_lists.append(tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=10))
callback_lists.append(tf.keras.callbacks.TensorBoard(log_dir="tensorboard/{}".format(settings), 
                                                     histogram_freq=0, 
                                                     batch_size=batch_size, 
                                                     write_graph=False, 
                                                     write_grads=False, 
                                                     write_images=False, 
                                                     embeddings_freq=0, 
                                                     embeddings_layer_names=None, 
                                                     embeddings_metadata=None, 
                                                     embeddings_data=None, 
                                                     update_freq='epoch'
                                                     ))
callback_lists.append(ModelCheckpoint(os.path.join("checkpoints/{}".format(settings), 'gru2gru_{epoch:03d}.hdf5'),
                                   verbose=1, save_weights_only=False, period=1))
# 开始训练
gru2gru_model.fit((train_x, np.zeros_like(train_y)),
                   train_y, 
                   batch_size=batch_size, 
                   validation_data=((test_x, np.zeros_like(test_y)), 
                   test_y), 
                   epochs=epochs,
                   callbacks=callback_lists)

# 保存模型
tf.saved_model.save(gru2gru_model, "saved_models/gru2gru/")